package cn.yx.application.gateway.subscriber;

import cn.hutool.core.bean.BeanUtil;
import cn.yx.common.core.constant.MdcConstant;
import org.reactivestreams.Subscription;
import org.slf4j.MDC;
import reactor.core.CoreSubscriber;
import reactor.util.context.Context;

import java.util.Optional;

/**
 * <p>MDC订阅器，实现从SkyWalking上下文中获取traceId并添加到MDC上下文</p>
 *
 * @param <T>
 * @author Wgssmart
 */
public class MdcSkyWalkingSubscriber<T> implements CoreSubscriber<T> {

    private static final String TRACE_ID = "traceId";

    private static final String SKY_WALKING_CTX_SNAPSHOT = "SKYWALKING_CONTEXT_SNAPSHOT";

    private final CoreSubscriber<T> actual;

    public MdcSkyWalkingSubscriber(CoreSubscriber<T> actual) {
        this.actual = actual;
    }

    @Override
    public void onSubscribe(Subscription subscription) {
        actual.onSubscribe(subscription);
    }

    @Override
    public void onNext(T t) {
        // 获取下一个数据流的上下文信息，并判断当前上下文中是否有SkyWalking的上下文信息，如果有，则从SkyWalking的上下文信息中获取traceId
        Context context = actual.currentContext();
        if (!context.isEmpty() && context.hasKey(SKY_WALKING_CTX_SNAPSHOT)) {
            Optional<String> traceIdOptional = Optional.of(context.get(SKY_WALKING_CTX_SNAPSHOT)).map(BeanUtil::beanToMap)
                    .map(o -> o.get(TRACE_ID)).map(BeanUtil::beanToMap).map(o -> o.get("id")).map(Object::toString);
            MDC.put(MdcConstant.SKY_WALKING_TID, traceIdOptional.orElse("TID:N/A"));
        }
        actual.onNext(t);
    }

    @Override
    public void onError(Throwable throwable) {
        actual.onError(throwable);
    }

    @Override
    public void onComplete() {
        actual.onComplete();
    }

    @Override
    public Context currentContext() {
        return actual.currentContext();
    }

}